from langchain_core.tools import tool
from langchain_core.prompts import ChatPromptTemplate,MessagesPlaceholder
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from model_utils import getLLM
from app_utils import websearch,repl
from typing import Annotated
"""
画出英国过去5年的国内生产总值的折线图
1.构建python 执行代码 PythonREPL工具
2.Painter类
1)初始Painter  __init__ 预构建带工具的_agent，生成_chain
2)__call__方法   调用 _chain.invoke方法调用python_repl工具， 画出英国过去5年的国内生产总值的折线图
3.main函数调用测试
"""

@tool
def python_repl(code:Annotated[str,"生成图表所需执行的python代码，图表用英文展示。"]):
    "使用它来执行python代码。如果你想看到一个值的输出，你应该用'print(...)'打印出来。这对用户是可见的"
    # print(code,code)
    try:
        result = repl.run(code)
    except BaseException as e:
        return f"Failed to execute.Error{repl(e)}"
    result_str = f"Successfully executed:\n```python\n{code}\n```\nStdout: {result}"
    return  (
        result_str + "\n\nIf you have completed all tasks, respond with FINAL ANSWER."
    )

_painner_system_template = '''
你是一个优秀的助手
根据提供工具回答问答，在你的回答前加上FINAL ANSWER，这样团队就知道该停下来了。
工具：{tools_name}
'''
class Painter:
    def __init__(self,llm):
        _tools = [python_repl]
        #消息站位
        _prompt = ChatPromptTemplate.from_messages([
            ("system",_painner_system_template),
            MessagesPlaceholder(variable_name="messages")
        ])

        _prompt = _prompt.partial(tools_name=",".join([_tool.name for _tool in _tools]))

        _agent = create_react_agent(llm,_tools)
        self._chain = _prompt | _agent

    def __call__(self,state):
        print("painter node state",state)
        _rt = self._chain.invoke(state)
        _messages = _rt["messages"]
        return _messages[-1]

if __name__ == '__main__':
    _llm = getLLM()
    _painter = Painter(_llm)
    _query = "获取英国过去5年的国内生产总值，然后画一个折线图，一旦你把它编码好，就完成。"
    # _query = "获取英国过去5年的国内生产总值，然后画一个折线图，一旦你把它编码好，并执行该代码，保存到当前路径下 名字为1.png。"
    _rt = _painter({"messages":["human",_query]})
    print(_rt)